// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/prestosql/aggregates/central_moments_aggregates.h>
#include <pollux/exec/aggregate.h>
#include <pollux/functions/lib/aggregates/central_moments_aggregates_base.h>
#include <pollux/functions/prestosql/aggregates/aggregate_names.h>

using namespace kumo::pollux::functions::aggregate;

namespace kumo::pollux::aggregate::prestosql {
    struct SkewnessResultAccessor {
        static bool hasResult(const CentralMomentsAccumulator &accumulator) {
            return accumulator.count() >= 3;
        }

        static double result(const CentralMomentsAccumulator &accumulator) {
            return std::sqrt(accumulator.count()) * accumulator.m3() /
                   std::pow(accumulator.m2(), 1.5);
        }
    };

    struct KurtosisResultAccessor {
        static bool hasResult(const CentralMomentsAccumulator &accumulator) {
            return accumulator.count() >= 4;
        }

        static double result(const CentralMomentsAccumulator &accumulator) {
            double count = accumulator.count();
            double m2 = accumulator.m2();
            double m4 = accumulator.m4();
            return ((count - 1) * count * (count + 1)) / ((count - 2) * (count - 3)) *
                   m4 / (m2 * m2) -
                   3 * ((count - 1) * (count - 1)) / ((count - 2) * (count - 3));
        }
    };

    template<typename TResultAccessor>
    exec::AggregateRegistrationResult registerCentralMoments(
        const std::string &name,
        bool withCompanionFunctions,
        bool overwrite) {
        std::vector<std::shared_ptr<exec::AggregateFunctionSignature> > signatures;
        std::vector<std::string> inputTypes = {
            "smallint", "integer", "bigint", "real", "double"
        };
        for (const auto &inputType: inputTypes) {
            signatures.push_back(
                exec::AggregateFunctionSignatureBuilder()
                .returnType("double")
                .intermediateType(CentralMomentsIntermediateResult::type())
                .argumentType(inputType)
                .build());
        }

        return exec::registerAggregateFunction(
            name,
            std::move(signatures),
            [name](
        core::AggregationNode::Step step,
        const std::vector<TypePtr> &argTypes,
        const TypePtr &resultType,
        const core::QueryConfig & /*config*/)
        -> std::unique_ptr<exec::Aggregate> {
                POLLUX_CHECK_LE(
                    argTypes.size(), 1, "{} takes at most one argument", name);
                const auto &inputType = argTypes[0];
                if (exec::isRawInput(step)) {
                    switch (inputType->kind()) {
                        case TypeKind::SMALLINT:
                            return std::make_unique<
                                CentralMomentsAggregatesBase<int16_t, TResultAccessor> >(
                                resultType);
                        case TypeKind::INTEGER:
                            return std::make_unique<
                                CentralMomentsAggregatesBase<int32_t, TResultAccessor> >(
                                resultType);
                        case TypeKind::BIGINT:
                            return std::make_unique<
                                CentralMomentsAggregatesBase<int64_t, TResultAccessor> >(
                                resultType);
                        case TypeKind::DOUBLE:
                            return std::make_unique<
                                CentralMomentsAggregatesBase<double, TResultAccessor> >(
                                resultType);
                        case TypeKind::REAL:
                            return std::make_unique<
                                CentralMomentsAggregatesBase<float, TResultAccessor> >(
                                resultType);
                        default:
                            POLLUX_UNSUPPORTED(
                                "Unsupported input type: {}. "
                                "Expected SMALLINT, INTEGER, BIGINT, DOUBLE or REAL.",
                                inputType->toString());
                    }
                } else {
                    checkAccumulatorRowType(
                        inputType,
                        "Input type for final aggregation must be "
                        "(count:bigint, m1:double, m2:double, m3:double, m4:double) struct");
                    return std::make_unique<CentralMomentsAggregatesBase<
                        int64_t /*unused*/,
                        TResultAccessor> >(resultType);
                }
            },
            withCompanionFunctions,
            overwrite);
    }

    void registerCentralMomentsAggregates(
        const std::string &prefix,
        bool withCompanionFunctions,
        bool overwrite) {
        registerCentralMoments<KurtosisResultAccessor>(
            prefix + kKurtosis, withCompanionFunctions, overwrite);
        registerCentralMoments<SkewnessResultAccessor>(
            prefix + kSkewness, withCompanionFunctions, overwrite);
    }
} // namespace kumo::pollux::aggregate::prestosql
